"""
异步任务执行器
专门处理轻度任务（IO 密集型）
"""

import asyncio
from typing import Any, Optional
from loguru import logger
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, TimeoutError

from app.base.base_task_scheduler import BaseTaskExecutor
from app.core.task_scheduler.task_monitor import task_monitor
from .progress_tracker import ProgressTracker
from .task_registry import task_registry
from app.models.task.task_data import Task


class AsyncExecutor(BaseTaskExecutor):
    """异步任务执行器"""

    def __init__(self):
        super().__init__("AsyncExecutor")

    async def execute_task(self, task: Task) -> Any:
        """执行异步任务"""
        logger.info(f"AsyncExecutor 开始执行任务: {task.id}")

        try:
            # 准备进度上下文
            progress_ctx = ProgressTracker(task.id, task.name)

            # 获取任务信息
            task_info = task_registry.get_task_info(task.func_name)

            try:
                # 准备参数
                kwargs = task.kwargs.copy()
                kwargs["task_id"] = task.id
                kwargs["task_name"] = task.name
                kwargs["progress_ctx"] = progress_ctx

                # 获取并执行任务函数
                task_func = task_registry.get_function(task.func_name)
                if not task_func:
                    raise ValueError(f"未找到任务函数: {task.func_name}")

                logger.debug(
                    f"执行任务函数: {task.func_name}, 参数: {task.args}, {kwargs}"
                )

                # 执行任务
                if asyncio.iscoroutinefunction(task_func):
                    result = await task_func(*task.args, **kwargs)
                else:
                    result = task_func(*task.args, **kwargs)

                logger.info(f"任务 {task.id} 执行成功")
                return result

            except Exception as e:
                logger.error(f"任务 {task.id} 执行失败: {e}")
                # 标记进度失败
                try:
                    await progress_ctx.fail(str(e))
                except Exception:
                    pass
                raise e

        except Exception as e:
            logger.error(f"AsyncExecutor 执行任务 {task.id} 时发生错误: {e}")
            raise e


# 创建全局实例
async_executor = AsyncExecutor()


def execute_task_in_process(
    func_module: str,
    func_name: str,
    args: tuple,
    kwargs: dict,
    task_id: Optional[str] = None,
    task_name: Optional[str] = None,
):
    """在独立进程中执行任务的函数"""
    try:
        # 动态导入函数所在模块
        module = __import__(func_module, fromlist=[func_name])
        func = getattr(module, func_name)

        # 创建简单的进度回调（多进程环境中的进度追踪比较复杂）
        # 这里只提供基础支持，实际项目中可能需要使用队列或共享内存
        if task_id and task_name:
            kwargs["task_id"] = task_id
            kwargs["task_name"] = task_name

        # 执行函数
        result = func(*args, **kwargs)
        return result
    except Exception as e:
        raise e


class ProcessTaskExecutor(BaseTaskExecutor):
    """多进程任务执行器 - 用于重度任务"""

    def __init__(self, max_workers: Optional[int] = None):
        super().__init__("ProcessTaskExecutor")
        self.max_workers = max_workers or min(4, mp.cpu_count())
        self._executor: Optional[ProcessPoolExecutor] = None
        self._lock = asyncio.Lock()

    async def _ensure_executor(self):
        """确保进程池执行器已创建"""
        if self._executor is None:
            async with self._lock:
                if self._executor is None:
                    self._executor = ProcessPoolExecutor(max_workers=self.max_workers)
                    logger.info(f"创建进程池执行器，最大工作进程数: {self.max_workers}")

    async def execute_task(self, task: Task) -> Any:
        """执行多进程任务"""
        await self._ensure_executor()

        if not self._executor:
            raise RuntimeError("进程池执行器未初始化")

        # 获取函数信息
        task_info = task_registry.get_task_info(task.func_name)
        timeout = task_info.get("timeout") or task.timeout

        # 检查函数是否适合多进程执行
        func = self.get_function(task.func_name)
        if asyncio.iscoroutinefunction(func):
            raise ValueError(
                f"协程函数 {task.func_name} 不能在进程池中执行，请使用AsyncExecutor"
            )

        # 启动监控
        if timeout:
            await task_monitor.start_monitoring(task, timeout)

        try:
            # 获取函数所在模块
            func_module = func.__module__

            # 在进程池中执行任务，传递任务ID和名称
            loop = asyncio.get_event_loop()
            future = self._executor.submit(
                execute_task_in_process,
                func_module,
                task.func_name,
                task.args,
                task.kwargs,
                task.id,
                task.name,
            )

            # 等待执行完成
            try:
                result = await asyncio.wait_for(
                    loop.run_in_executor(None, future.result), timeout=timeout
                )
                return result
            except asyncio.TimeoutError:
                # 取消任务
                future.cancel()
                raise TimeoutError(f"任务执行超时: {timeout}秒")

        finally:
            # 停止监控
            await task_monitor.stop_monitoring(task.id)

    async def shutdown(self):
        """关闭进程池执行器"""
        logger.info("正在关闭进程池执行器...")
        self.stop()  # 标记为停止状态

        if self._executor:
            async with self._lock:
                if self._executor:
                    try:
                        # 设置较短的等待时间，避免无限等待
                        self._executor.shutdown(wait=False)

                        # 给一些时间让正在执行的任务完成
                        import time

                        start_time = time.time()
                        timeout = 5.0  # 5秒超时

                        while (
                            self._executor._processes
                            and (time.time() - start_time) < timeout
                        ):
                            await asyncio.sleep(0.1)

                        # 如果还有进程在运行，强制终止
                        if (
                            hasattr(self._executor, "_processes")
                            and self._executor._processes
                        ):
                            logger.warning("强制终止剩余的工作进程")

                        self._executor = None
                        logger.info("进程池执行器已关闭")

                    except Exception as e:
                        logger.error(f"关闭进程池执行器时出错: {e}")
                        self._executor = None

    def get_stats(self) -> dict:
        """获取执行器统计信息"""
        if not self._executor:
            return {"status": "未初始化", "max_workers": self.max_workers}

        return {
            "status": "运行中",
            "max_workers": self.max_workers,
            "is_running": self.is_running,
        }


process_executor = ProcessTaskExecutor()
